import os
import tkinter as tk
from tkinter import filedialog, messagebox, ttk
import hashlib
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.backends import default_backend
import base64
import sys

# 确保中文显示正常
if sys.platform.startswith('win'):
    import matplotlib
    matplotlib.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]

class AESEncryptorApp:
    def __init__(self, root):
        self.root = root
        self.root.title("文件加密解密工具")
        self.root.geometry("600x400")
        self.root.resizable(True, True)
        
        # 设置主题颜色
        self.bg_color = "#f0f0f0"
        self.button_color = "#4CAF50"
        self.button_text_color = "white"
        self.root.configure(bg=self.bg_color)
        
        # 创建主框架
        self.main_frame = ttk.Frame(root, padding="20")
        self.main_frame.pack(fill=tk.BOTH, expand=True)
        
        # 文件路径变量
        self.file_path = tk.StringVar()
        self.custom_extension = tk.StringVar(value="abc")  # 默认扩展名
        
        # 创建UI元素
        self.create_widgets()
        
    def create_widgets(self):
        # 标题
        title_label = ttk.Label(
            self.main_frame, 
            text="AES-256 文件加密解密工具", 
            font=("SimHei", 16, "bold")
        )
        title_label.grid(row=0, column=0, columnspan=3, pady=10)
        
        # 文件路径输入框
        ttk.Label(self.main_frame, text="文件路径:", font=("SimHei", 10)).grid(row=1, column=0, sticky=tk.W, pady=5)
        file_entry = ttk.Entry(self.main_frame, textvariable=self.file_path, width=50)
        file_entry.grid(row=1, column=1, pady=5)
        
        # 浏览按钮
        browse_btn = ttk.Button(
            self.main_frame, 
            text="浏览", 
            command=self.browse_file
        )
        browse_btn.grid(row=1, column=2, padx=5, pady=5)
        
        # 自定义扩展名
        ttk.Label(self.main_frame, text="加密后扩展名:", font=("SimHei", 10)).grid(row=2, column=0, sticky=tk.W, pady=5)
        ext_entry = ttk.Entry(self.main_frame, textvariable=self.custom_extension, width=10)
        ext_entry.grid(row=2, column=1, sticky=tk.W, pady=5)
        ttk.Label(self.main_frame, text="(无需输入.)", font=("SimHei", 8)).grid(row=2, column=2, sticky=tk.W)
        
        # 密码输入
        ttk.Label(self.main_frame, text="密码:", font=("SimHei", 10)).grid(row=3, column=0, sticky=tk.W, pady=5)
        self.password_entry = ttk.Entry(self.main_frame, show="*", width=30)
        self.password_entry.grid(row=3, column=1, sticky=tk.W, pady=5)
        
        # 按钮框架
        button_frame = ttk.Frame(self.main_frame)
        button_frame.grid(row=4, column=0, columnspan=3, pady=20)
        
        # 加密按钮
        encrypt_btn = ttk.Button(
            button_frame, 
            text="加密文件", 
            command=self.encrypt_file,
            width=15
        )
        encrypt_btn.pack(side=tk.LEFT, padx=10)
        
        # 解密按钮
        decrypt_btn = ttk.Button(
            button_frame, 
            text="解密文件", 
            command=self.decrypt_file,
            width=15
        )
        decrypt_btn.pack(side=tk.LEFT, padx=10)
        
        # 状态框
        ttk.Label(self.main_frame, text="状态:", font=("SimHei", 10)).grid(row=5, column=0, sticky=tk.NW, pady=5)
        self.status_text = tk.Text(self.main_frame, height=8, width=50, wrap=tk.WORD)
        self.status_text.grid(row=5, column=1, columnspan=2, pady=5)
        self.status_text.config(state=tk.DISABLED)
        
        # 底部信息
        footer_label = ttk.Label(
            self.main_frame, 
            text="支持加密RAR等各类文件，加密后需用本工具解密", 
            font=("SimHei", 8)
        )
        footer_label.grid(row=6, column=0, columnspan=3, pady=10)
        
    def browse_file(self):
        """浏览并选择文件"""
        file_path = filedialog.askopenfilename()
        if file_path:
            self.file_path.set(file_path)
            self.update_status(f"已选择文件: {os.path.basename(file_path)}")
    
    def update_status(self, message):
        """更新状态文本框"""
        self.status_text.config(state=tk.NORMAL)
        self.status_text.insert(tk.END, message + "\n")
        self.status_text.see(tk.END)
        self.status_text.config(state=tk.DISABLED)
        self.root.update_idletasks()
    
    def get_key(self, password):
        """从密码生成AES密钥"""
        password_bytes = password.encode()
        salt = b'salt_for_aes_encryption'  # 固定盐值，实际应用中可随机生成并存储
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
            backend=default_backend()
        )
        key = base64.urlsafe_b64encode(kdf.derive(password_bytes))
        return key
    
    def encrypt_file(self):
        """加密文件"""
        file_path = self.file_path.get()
        password = self.password_entry.get()
        ext = self.custom_extension.get().strip()
        
        # 验证输入
        if not file_path:
            messagebox.showerror("错误", "请选择要加密的文件")
            return
        
        if not password:
            messagebox.showerror("错误", "请输入密码")
            return
        
        if not ext:
            messagebox.showerror("错误", "请输入自定义扩展名")
            return
        
        try:
            self.update_status(f"开始加密文件: {os.path.basename(file_path)}")
            
            # 生成密钥
            key = self.get_key(password)
            fernet = Fernet(key)
            
            # 读取文件内容
            with open(file_path, 'rb') as file:
                data = file.read()
            
            # 加密数据
            encrypted_data = fernet.encrypt(data)
            
            # 构建新文件名
            new_file_path = f"{file_path}.{ext}"
            
            # 写入加密文件
            with open(new_file_path, 'wb') as file:
                file.write(encrypted_data)
            
            self.update_status(f"文件加密成功!")
            self.update_status(f"加密文件保存为: {os.path.basename(new_file_path)}")
            messagebox.showinfo("成功", f"文件加密成功!\n保存为: {new_file_path}")
            
        except Exception as e:
            error_msg = f"加密失败: {str(e)}"
            self.update_status(error_msg)
            messagebox.showerror("错误", error_msg)
    
    def decrypt_file(self):
        """解密文件"""
        file_path = self.file_path.get()
        password = self.password_entry.get()
        ext = self.custom_extension.get().strip()
        
        # 验证输入
        if not file_path:
            messagebox.showerror("错误", "请选择要解密的文件")
            return
        
        if not password:
            messagebox.showerror("错误", "请输入密码")
            return
        
        # 检查文件扩展名是否匹配
        if not file_path.endswith(f".{ext}"):
            if not messagebox.askyesno("警告", f"所选文件似乎不是.{ext}格式的加密文件，仍要尝试解密吗?"):
                return
        
        try:
            self.update_status(f"开始解密文件: {os.path.basename(file_path)}")
            
            # 生成密钥
            key = self.get_key(password)
            fernet = Fernet(key)
            
            # 读取加密文件内容
            with open(file_path, 'rb') as file:
                encrypted_data = file.read()
            
            # 解密数据
            decrypted_data = fernet.decrypt(encrypted_data)
            
            # 构建新文件名（移除自定义扩展名）
            if file_path.endswith(f".{ext}"):
                new_file_path = file_path[:-(len(ext)+1)]
                # 如果文件已存在，添加后缀
                if os.path.exists(new_file_path):
                    name, ext_orig = os.path.splitext(new_file_path)
                    new_file_path = f"{name}_decrypted{ext_orig}"
            else:
                name, ext_orig = os.path.splitext(file_path)
                new_file_path = f"{name}_decrypted{ext_orig}"
            
            # 写入解密文件
            with open(new_file_path, 'wb') as file:
                file.write(decrypted_data)
            
            self.update_status(f"文件解密成功!")
            self.update_status(f"解密文件保存为: {os.path.basename(new_file_path)}")
            messagebox.showinfo("成功", f"文件解密成功!\n保存为: {new_file_path}")
            
        except Exception as e:
            error_msg = f"解密失败: {str(e)}\n可能是密码错误或文件已损坏"
            self.update_status(error_msg)
            messagebox.showerror("错误", error_msg)

if __name__ == "__main__":
    root = tk.Tk()
    app = AESEncryptorApp(root)
    root.mainloop()
